'''
Author: LIFM0623 1944099171@qq.com
Date: 2024-02-21 21:42:12
LastEditors: LIFM0623 1944099171@qq.com
LastEditTime: 2024-03-13 14:27:50
FilePath: \scene_back_nest\src\imageraster\python\tif2png.py
Description:  目前只可进行 单波段 与 多波段进行转换
'''
import sys
from osgeo import gdal,osr
import os
import matplotlib
import numpy as np
from matplotlib import pyplot as plt

# 获取tif的bbox 边界
def get_bbox(dataset):
    # 获取数据集的投影
    projection = dataset.GetProjection()
    # 检查投影是否为地理坐标系（0/1）
    spatialRef = osr.SpatialReference(projection)
    isGeographic = spatialRef.IsGeographic()
    # 获取外边框信息
    geoTransform = dataset.GetGeoTransform()
    topLeftX = geoTransform[0]
    topLeftY = geoTransform[3]
    pixelWidth = geoTransform[1]
    pixelHeight = geoTransform[5]
    width = dataset.RasterXSize
    height = dataset.RasterYSize
    # 计算左下角经纬度
    bottomLeftX = topLeftX
    bottomLeftY = topLeftY + (height * pixelHeight)
    # 计算右上角经纬度
    topRightX = topLeftX + (width * pixelWidth)
    topRightY = topLeftY
    # 如果投影是地理坐标系，则直接使用地理坐标
    if isGeographic == 1:
        bottomLeft = [bottomLeftX, bottomLeftY]
        topRight = [topRightX, topRightY]
    else:
        # 否则，进行坐标转换
        # 创建地理坐标转换器
        geographicRef = spatialRef.CloneGeogCS()
        transform = osr.CoordinateTransformation(spatialRef, geographicRef)
        # 转换为经纬度坐标
        bottomLeft = transform.TransformPoint(bottomLeftX, bottomLeftY)
        topRight = transform.TransformPoint(topRightX, topRightY)
        # 平面坐标系的（x,y)分别对应地理坐标系的（维度，经度）,所以需要调换一下顺序
        bottomLeft = [bottomLeft[1], bottomLeft[0]]
        topRight = [topRight[1], topRight[0]]

    # 存储结果到List[float]
    bbox = []
    bbox.append(bottomLeft[0])  # 左下经度
    bbox.append(bottomLeft[1])  # 左下纬度
    bbox.append(topRight[0])  # 右上经度
    bbox.append(topRight[1])  # 右上纬度
    bboxString = str(bottomLeft[0]) +','+ str(bottomLeft[1]) +','+ str(topRight[0]) +','+ str(topRight[1])

    return bboxString


# 彩色波段 tifTopng
def RGB_tif_to_PNG(input_path:str,output_path:str):
    # 读取数据
    src_ds = gdal.Open(input_path)
 
    # 设置转换选项
    options = gdal.TranslateOptions(format='PNG', bandList=[3, 2, 1])
    # 构建输出路径
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
 
    # 保存转换数据
    gdal.Translate(output_path, src_ds, options=options)
 
    # 关闭数据源
    src_ds = None

# 单波段 tifTopng
def Single_tif_to_Png(input_path:str,output_path:str):
    dataset = gdal.Open(input_path)
    band = dataset.GetRasterBand(1)
    # 提取NoData值
    no_data_value = band.GetNoDataValue()
    # print("NoData值为:", no_data_value)

    # 读取数组
    elevation_data = band.ReadAsArray()
    # 将无效值替换为nan，防止渲染时被渲染
    # elevation_data[elevation_data == -32768] = np.nan 这种方法兼容性较差，容易出现类型转换错误
    if no_data_value is not None:
        elevation_data = np.where(
            elevation_data == no_data_value, np.nan, elevation_data
        )  # 这种方法似乎会将整型的数组像元值转成float类型，但不影响其他
    # 将0值替换为nan，避免被渲染（一是避免无效值没有写明，而是避免无效值不等于0）
    elevation_data = np.where(elevation_data == 0, np.nan, elevation_data)

    # print("Min value: ", np.nanmin(elevation_data))
    # print("Max value: ", np.nanmax(elevation_data))
    # color_map = plt.cm.get_cmap("jet")
    color_map = matplotlib.colormaps["jet"]
    # color_map.set_under(alpha=0)  # 将像元值低于vmin以下的的颜色设为透明

    plt.imshow(elevation_data, cmap=color_map)
    # plt.colorbar()
    plt.axis("off")
    # 注：  plt.show() 与 plt.savefig() 不可以同时使用  否则savefig 的图片无颜色
    # plt.show()
    # 保存图像为PNG文件
    plt.savefig(
        output_path,
        dpi=300,
        bbox_inches="tight",
        pad_inches=0,
        transparent=True,
    )

    dataset = None

if __name__ == "__main__":
    # 获取关键字参数
    file_path = sys.argv[1]
    output_path = sys.argv[2]

    dataset = gdal.Open(file_path)

    bbox = get_bbox(dataset)
    print(bbox)
    band_count = dataset.RasterCount
    if(band_count == 1):
        Single_tif_to_Png(file_path,output_path)
    else:
        print('Error')
        RGB_tif_to_PNG(file_path,output_path)